package com.mask.encryption.service.impl;

import com.mask.encryption.service.MaskEncryptionService;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA-OAEP(SHA-256) 实现
 * 注意：RSA 适合小数据块；如需大数据加密建议采用混合加密（RSA加密AES密钥，数据用AES加密）。
 */
public class RsaOaepEncryptionService implements MaskEncryptionService {

    private static final String TRANSFORMATION = "RSA/ECB/OAEPWithSHA-256AndMGF1Padding";

    private final PublicKey publicKey;
    private final PrivateKey privateKey;

    public RsaOaepEncryptionService(String publicKeyPem, String privateKeyPem) {
        try {
            this.publicKey = parsePublicKey(publicKeyPem);
            this.privateKey = parsePrivateKey(privateKeyPem);
        } catch (Exception e) {
            throw new IllegalArgumentException("Invalid RSA keys", e);
        }
    }

    private static PublicKey parsePublicKey(String pem) throws Exception {
        String normalized = pem
                .replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s", "");
        byte[] keyBytes = Base64.getDecoder().decode(normalized);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
        return KeyFactory.getInstance("RSA").generatePublic(spec);
    }

    private static PrivateKey parsePrivateKey(String pem) throws Exception {
        String normalized = pem
                .replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s", "");
        byte[] keyBytes = Base64.getDecoder().decode(normalized);
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
        return KeyFactory.getInstance("RSA").generatePrivate(spec);
    }

    @Override
    public byte[] encrypt(byte[] plaintext) {
        try {
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            return cipher.doFinal(plaintext);
        } catch (Exception e) {
            throw new IllegalStateException("RSA-OAEP encrypt failed", e);
        }
    }

    @Override
    public byte[] decrypt(byte[] ciphertext) {
        try {
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            return cipher.doFinal(ciphertext);
        } catch (Exception e) {
            throw new IllegalStateException("RSA-OAEP decrypt failed", e);
        }
    }

    @Override
    public String encryptToBase64(String plaintext) {
        return Base64.getEncoder().encodeToString(encrypt(plaintext.getBytes(StandardCharsets.UTF_8)));
    }

    @Override
    public String decryptFromBase64(String base64Ciphertext) {
        byte[] cipher = Base64.getDecoder().decode(base64Ciphertext);
        return new String(decrypt(cipher), StandardCharsets.UTF_8);
    }
}


